Ground state search for a Rydberg system (Rb87)

Rb87 with Rydberg state 70s1/2) in a 3d setup.

# pylint: disable=invalid-name

import numpy as np
import numpy.linalg as nla

import qtealeaves as qtl
from qtealeaves import modeling


def main(tn_type=5, input_folder=None, output_folder=None):
    """
    Main method for the ground state search of the 3d Rydberg model.

    **Arguments**

    tn_type : int, optional
        Choose 5 for python-TTN, 6 for python-MPS.
        Default to 5.

    input_folder : str | None, optional
        Input folder. Default to None.

    output_folder : str | None, optional
        Output folder. Default to None.
    """
    model_name = lambda params: "Rydberg_L%d" % (params["L"])
    if input_folder is None:
        input_folder = lambda params: "Rb87/input_L%d" % (params["L"])
    if output_folder is None:
        output_folder = lambda params: "Rb87/output_L%d" % (params["L"])

    model = modeling.QuantumModel(3, "L", name=model_name)

    model += modeling.LocalTerm("sx", strength="Jx")
    model += modeling.LocalTerm("n", strength="U")

    model += modeling.TwoBodyTerm3D(["n", "n"], [1, 0, 0], strength="V1")
    model += modeling.TwoBodyTerm3D(["n", "n"], [1, 1, 0], strength="V2")
    model += modeling.TwoBodyTerm3D(["n", "n"], [1, 1, 1], strength="V3")
    model += modeling.TwoBodyTerm3D(["n", "n"], [2, 0, 0], strength="V4")

    my_conv = qtl.convergence_parameters.TNConvergenceParameters(
        max_iter=3, max_bond_dimension=2
    )
    my_ops = qtl.operators.TNSpin12Operators()
    my_obs = qtl.observables.TNObservables()
    my_obs += qtl.observables.TNObsLocal("<sx>", "sx")

    simulation = qtl.QuantumGreenTeaSimulation(
        model,
        my_ops,
        my_conv,
        my_obs,
        tn_type=tn_type,
        folder_name_input=input_folder,
        folder_name_output=output_folder,
        store_checkpoints=False,
    )

    params = [
        {"L": 4, "Jx": 2.0, "U": -15.0, "V1": 25.0, "V2": 3.12, "V3": 0.92, "V4": 0.39}
    ]

    simulation.run(params, delete_existing_folder=True)

    for elem in params:
        tn_energy_0 = simulation.get_static_obs(elem)["energy"]

        if elem["L"] == 2:
            ed_energy_0 = nla.eigh(model.build_ham(my_ops, elem))[0][0]
            print("TN vs ED ground state energy", tn_energy_0, ed_energy_0)

            assert np.abs(tn_energy_0 - ed_energy_0) < 1e-4
            msg_check = "ground state energy at least correct up to 1e-4."

        else:
            print("TN ground state energy", tn_energy_0)
            msg_check = "no asserts implemented for this system size."

    print(f"\nExample `{__file__}` ran successfully; {msg_check}")

    return


if __name__ == "__main__":
    main()